from __future__ import print_function

from numba import jit
import os.path
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from skimage import io
from sklearn.utils.linear_assignment_ import linear_assignment
import glob
import time
import argparse
from filterpy.kalman import KalmanFilter

# 用于修复有些结果宽度太窄
B_REQUEST_RESIZE = True


MODE_ASSOCIATION_IOU = 0
MODE_ASSOCIATION_JOINT_DIST = 1

class Config(object):
    F_PERFECT_BOX_RATIO = 2.1
    MODE_ASSOCIATION = MODE_ASSOCIATION_IOU

class KalmanJointTracker(object):
    """
  This class represents the internel state of individual tracked objects observed as bbox.
  """
    count = 0

    def __init__(self, joint_num, joints: np.ndarray, full_box):
        """
    Initialises a tracker using initial bounding box.
        :type joints: np.array
    """
        # define constant velocity model
        dim_x = joint_num * 4
        dim_z = joint_num * 2
        self.joint_num = joint_num
        self.kf = KalmanFilter(dim_x=dim_x, dim_z=dim_z)
        self.kf.F = np.eye(dim_x)
        self.kf.F[:dim_z, dim_z:] = np.eye(dim_z)
        self.kf.H = np.zeros((dim_z, dim_x))
        self.kf.H[:dim_z, :dim_z] = np.eye(dim_z)
        # self.kf.P[dim_z:, dim_z:] *= 1000.  # give high uncertainty to the unobservable initial velocities
        # self.kf.R[2:, 2:] *= 10.
        # self.kf.Q[-1, -1] *= 0.01
        self.kf.Q[dim_z:, dim_z:] *= 0.01

        # self.kf.P = np.eye(dim_x) * 10000.
        self.kf.P[dim_z:] *= 1000.
        self.kf.P *= np.eye(dim_x) * 10.
        # self.kf.R = np.eye(dim_z) * 10
        self.kf.R[dim_z:, dim_z:] *= 100
        self.kf.R[:dim_z, :dim_z] *= .1

        self.kf.x[:dim_z, 0] = joints.flatten()  # joints.reshape((-1, 1))
        self.time_since_update = 0
        self.id = KalmanJointTracker.count
        KalmanJointTracker.count += 1
        self.history = []
        self.hits = 0
        self.hit_streak = 0
        self.age = 0
        self.last_measure = None
        # self.init_box = full_box
        # self.init_in_box = self.convert_join2box(joints)

    def update(self, joints):
        """
    Updates the state vector with observed bbox.
    """
        self.time_since_update = 0
        self.history = []
        self.hits += 1
        self.hit_streak += 1
        messure_z = joints.reshape((-1, 1))
        self.kf.update(messure_z)
        self.last_measure = messure_z

    def predict(self):
        """
    Advances the state vector and returns the predicted bounding box estimate.
    """
        if (self.kf.x[6] + self.kf.x[2]) <= 0:  # ratio < 0
            self.kf.x[6] *= 0.0
        self.kf.predict()
        self.age += 1
        if self.time_since_update > 0:
            self.hit_streak = 0
        self.time_since_update += 1
        in_box = self.convert_join2box(self.kf.x)
        self.history.append(in_box)
        return self.history[-1]

    def get_state(self):
        """
        如果是预测的结果, 没有任何关联. 则使用上次关联对应的检测位置.
        Returns the current bounding box estimate.
        """
        # return self.convert_joint2box_by_estimate(self.kf.x)
        if self.last_measure is not None:
            state = self.convert_joint2box_by_estimate(self.last_measure)
        else:
            state = self.convert_joint2box_by_estimate(self.kf.x)
        self.last_measure = None

        # if self.last_measure is not None:
        #     state = self.convert_join2box(self.last_measure)
        # else:
        #     state = self.convert_join2box(self.kf.x)
        # self.last_measure = None
        return state

    def convert_join2box(self, x):
        joints = x[:self.joint_num * 2]
        points = joints.reshape((-1, 2))
        l, t = np.min(points[:, [0, 1]], axis=0)
        r, b = np.max(points[:, [0, 1]], axis=0)
        return np.array([l, t, r, b])

    def convert_joint2box_by_estimate(self, x):
        joints = x[:self.joint_num * 2]
        points = joints.reshape((-1, 2))
        # points = points[[2, 3, 6, 7, 8, 9, 12, 13], :]
        # points = points[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 12, 13], :]
        points = points[:, :]
        l, t = np.min(points[:, [0, 1]], axis=0)
        r, b = np.max(points[:, [0, 1]], axis=0)
        if B_REQUEST_RESIZE:
            height = b - t
            width = r - l
            min_width = height / Config.F_PERFECT_BOX_RATIO
            if width < min_width:
                centre_h = l + width / 2.
                l = centre_h - min_width / 2.
                r = centre_h + min_width / 2.

        return np.array([l, t, r, b])




@jit
def iou(bb_test, bb_gt):
    """
    Computes IUO between two bboxes in the form [x1,y1,x2,y2]
    """
    xx1 = np.maximum(bb_test[0], bb_gt[0])
    yy1 = np.maximum(bb_test[1], bb_gt[1])
    xx2 = np.minimum(bb_test[2], bb_gt[2])
    yy2 = np.minimum(bb_test[3], bb_gt[3])
    w = np.maximum(0., xx2 - xx1)
    h = np.maximum(0., yy2 - yy1)
    wh = w * h
    o = wh / ((bb_test[2] - bb_test[0]) * (bb_test[3] - bb_test[1])
              + (bb_gt[2] - bb_gt[0]) * (bb_gt[3] - bb_gt[1]) - wh)
    return (o)


def associate_detections_to_trackers(detections, trackers, iou_threshold=0.3):
    """
    Assigns detections to tracked object (both represented as bounding boxes)

    Returns 3 lists of matches, unmatched_detections and unmatched_trackers
    """
    if (len(trackers) == 0):
        return np.empty((0, 2), dtype=int), np.arange(len(detections)), np.empty((0, 5), dtype=int)
    iou_matrix = np.zeros((len(detections), len(trackers)), dtype=np.float32)

    for d, det in enumerate(detections):
        for t, trk in enumerate(trackers):
            iou_matrix[d, t] = iou(det, trk)
    matched_indices = linear_assignment(-iou_matrix)

    unmatched_detections = []
    for d, det in enumerate(detections):
        if (d not in matched_indices[:, 0]):
            unmatched_detections.append(d)
    unmatched_trackers = []
    for t, trk in enumerate(trackers):
        if (t not in matched_indices[:, 1]):
            unmatched_trackers.append(t)

    # filter out matched with low IOU
    matches = []
    for m in matched_indices:
        if (iou_matrix[m[0], m[1]] < iou_threshold):
            unmatched_detections.append(m[0])
            unmatched_trackers.append(m[1])
        else:
            matches.append(m.reshape(1, 2))
    if (len(matches) == 0):
        matches = np.empty((0, 2), dtype=int)
    else:
        matches = np.concatenate(matches, axis=0)

    return matches, np.array(unmatched_detections), np.array(unmatched_trackers)


from sort.util import get_parametric_distance


def associate_detections_to_trackers_by_joints(detections, trackers, iou_threshold=0.3):
    """
    Assigns detections to tracked object (both represented as bounding boxes)

    Returns 3 lists of matches, unmatched_detections and unmatched_trackers (id, box)
    """

    def joint_dist(det, trk):
        return 0

    if (len(trackers) == 0):
        return np.empty((0, 2), dtype=int), np.arange(len(detections)), np.empty((0, 5), dtype=int)
    dist_matrix = np.zeros((len(detections), len(trackers)), dtype=np.float32)

    for d, det in enumerate(detections):
        for t, trk in enumerate(trackers):
            dist_matrix[d, t] = get_parametric_distance(det, trk)
            # dist_matrix[d, t] = joint_dist(det, trk)
    # for d, det in enumerate(detections):
    #     dist_matrix[d, :] = get_parametric_distance(det, trackers)
    matched_indices = linear_assignment(dist_matrix)

    unmatched_detections = []
    for d, det in enumerate(detections):
        if d not in matched_indices[:, 0]:
            unmatched_detections.append(d)
    unmatched_trackers = []
    for t, trk in enumerate(trackers):
        if t not in matched_indices[:, 1]:
            unmatched_trackers.append(t)

    # filter out matched with low IOU
    matches = []
    # todo filter out
    for m in matched_indices:
        if dist_matrix[m[0], m[1]] >= 1000000.0:
            unmatched_detections.append(m[0])
            unmatched_trackers.append(m[1])
        else:
            matches.append(m.reshape(1, 2))
    if len(matches) == 0:
        matches = np.empty((0, 2), dtype=int)
    else:
        matches = np.concatenate(matches, axis=0)

    return matches, np.array(unmatched_detections), np.array(unmatched_trackers)


class Sort(object):
    def __init__(self, max_age=1, min_hits=3, iou_threshold=0.3):
        """
        Sets key parameters for SORT
        """
        self.max_age = max_age
        self.min_hits = min_hits
        self.trackers = []
        self.frame_count = 0
        self.iou_threshold = iou_threshold
        self.joint_num = 16  # todo

    def update(self, dets, joints):
        """
        Params:
          dets - a numpy array of detections in the format [[x1,y1,x2,y2,score],[x1,y1,x2,y2,score],...]
        Requires: this method must be called once for each frame even with empty detections.
        Returns the a similar array, where the last column is the object ID.
        joints: 关节点2d位置信息; (n, 16*2)

        NOTE: The number of objects returned may differ from the number of detections provided.
        """
        self.frame_count += 1
        # get predicted locations from existing trackers.
        trks = np.zeros((len(self.trackers), 5))
        ret = []

        # predict tracker by Kalman filter
        for t, trk in enumerate(trks):
            pos = self.trackers[t].predict()
            trk[:] = [pos[0], pos[1], pos[2], pos[3], 0]
        trks = np.ma.compress_rows(np.ma.masked_invalid(trks))
        if Config.MODE_ASSOCIATION == MODE_ASSOCIATION_JOINT_DIST:
            tck_joints = np.array([tracker.kf.x[:self.joint_num*2] for tracker in self.trackers])
            matched, unmatched_dets, unmatched_trks = associate_detections_to_trackers_by_joints(joints, tck_joints)
        else:
            matched, unmatched_dets, unmatched_trks = associate_detections_to_trackers(dets, trks, self.iou_threshold)
        matched_indexes = {}
        # update matched trackers with assigned detections
        for t, trk in enumerate(self.trackers):
            if (t not in unmatched_trks):
                d = matched[np.where(matched[:, 1] == t)[0], 0]     # matched detection index
                # trk.update(dets[d, :][0])
                trk.update(joints[d, :][0])
        # create and initialise new trackers for unmatched detections
        for i in unmatched_dets:
            trk = KalmanJointTracker(self.joint_num, joints[i, :], dets[i, :])
            self.trackers.append(trk)
        i = len(self.trackers)
        for trk in reversed(self.trackers):
            d = trk.get_state()
            if ((trk.time_since_update < 1) and (trk.hit_streak >= self.min_hits or self.frame_count <= self.min_hits)):
                ret.append(np.concatenate((d, [trk.id + 1])).reshape(1, -1))  # +1 as MOT benchmark requires positive
            i -= 1
            # remove dead tracklet
            if (trk.time_since_update > self.max_age):
                self.trackers.pop(i)
        if len(ret) > 0:
            return np.concatenate(ret)
        return np.empty((0, 5))


# test


def oneshot_15_pose(pose_df_path, output_path):
    import pandas as pd

    def get_joints(dataframes: pd.DataFrame):
        l = 16
        col_names = [j + str(i) for i in range(l) for j in ['x', 'y']]
        return dataframes.loc[:, col_names].values

    # mot_tracker = Sort(iou_threshold=0.1)
    # mot_tracker = Sort(max_age=15, min_hits=8, iou_threshold=0.3) # for eth-sunnyday
    # todo 这里没有进行调参
    mot_tracker = Sort(max_age=7, min_hits=2, iou_threshold=0.2)  # for pets s2l1

    df = pd.read_hdf(pose_df_path, 'df')

    with open(output_path, 'w') as out_file:
        time_stamp1 = time.time()
        for frame in np.unique(df.fid.values):
            # frame += 1
            # print(frame)
            cur_df = df[df.fid == frame]
            boxes = cur_df.box
            joints = get_joints(cur_df)
            # assert joints.shape[1] == 16 * 2  # todo magic number
            dets = np.vstack([boxes.iloc[i] for i in range(len(boxes))])
            dets[:, 2:4] += dets[:, 0:2]  # convert to [x1,y1,w,h] to [x1,y1,x2,y2]
            trackers = mot_tracker.update(dets, joints)
            for d in trackers:
                tlwh_box = [d[0], d[1], d[2] - d[0], d[3] - d[1]]
                tid = d[4]
                # dd = estimate_box(tlwh_box)
                dd = tlwh_box
                print('%d,%d,%.2f,%.2f,%.2f,%.2f,1,-1,-1,-1' % (frame, tid, dd[0], dd[1], dd[2], dd[3]),
                      file=out_file)
        t2 = time.time()
        print('fps : {}'.format(df.fid.values[-1] / (t2 - time_stamp1)))


def test():
    dataset_name = 'PETS09-S2L1'  # TUD-Stadtmitte
    # dataset_name = 'ETH-Sunnyday'
    # dataset_name = 'ETH-Pedcross2'
    pose_df = '../output/mot15_pose/train/{}/det_pose.hdf5'.format(dataset_name)
    # output_path = './test.txt'
    output_path = os.path.join('../output/tuning/', dataset_name)
    oneshot_15_pose(pose_df, output_path)


if __name__ == '__main__':
    test()
